#ifdef ENABLE_MAPLE_SAN

#include "asan_ud.h"


namespace maple {

UDProfile::UDProfile() {}

set_check *UDProfile::GetSetCheck(StmtID stmtID){
  // case for user check 
  if (this->StmtID2Check.find(stmtID) != this->StmtID2Check.end()){
    return &(this->StmtID2Check[stmtID]); 
  }
  // for asan 
  if (this->dassignID2dep.find(stmtID) != this->dassignID2dep.end()){
    return &(this->dassignID2dep[stmtID]); 
  }
  return nullptr; 
}

void UDProfile::CreateDassignIDMap() {
  for (auto& [stmtID, check] : this->StmtID2Check) {
    for (auto& dassignID : check.dassignID) {
      // Hitting this for 5 times in bzip2 
      // if (this->dassignID2dep.find(dassignID) != this->dassignID2dep.end()) {
      //   LogInfo::MapleLogger() << "Hitting a dup in CreateDassignIDMap "<<" \n";
      // }
      dassignID2dep[dassignID] = check; 
    }
  }
}


StmtNode* UDProfile::retLatest_Regassignment(StmtNode *stmt, RegID register_number) {
  /*
    Given a register number, return the stmt node 
  */
  StmtNode *ret_stmt = nullptr;
  StmtNode *prevStmt = stmt->GetPrev();
  if (prevStmt != nullptr) {
    if (prevStmt->GetOpCode() == OP_regassign) {
      RegassignNode *regAssign = static_cast<RegassignNode *>(prevStmt);
      if (register_number == regAssign->GetRegIdx()) {
        return prevStmt;
      } else {
        ret_stmt = retLatest_Regassignment(prevStmt, register_number);
      }
    } else if (prevStmt->GetOpCode() == OP_iassign) {
      IassignNode *iassign = static_cast<IassignNode *>(prevStmt);
      BaseNode *addr_expr = iassign->Opnd(0);
      if (addr_expr->GetOpCode() == OP_iread) {
        std::vector<VarID> dump_reg;
        std::vector<IntVal> dump_const;
        recursion(addr_expr, dump_reg, dump_const);
        for (auto reg_tmp : dump_reg) {
          if (reg_tmp == register_number) {
            return prevStmt;
          }
        }
        ret_stmt = retLatest_Regassignment(prevStmt, register_number);
      } else if (addr_expr->GetOpCode() == OP_regread) {
        RegreadNode *regread = static_cast<RegreadNode *>(addr_expr);
        if (register_number == regread->GetRegIdx()) {
          return prevStmt;
        } else {
          ret_stmt = retLatest_Regassignment(prevStmt, register_number);
        }
      } else if (IsCommutative(addr_expr->GetOpCode())) {
        /*
          0th stmt: add u64 (
            iread u64 <* <$_TY_IDX111>> 22 (regread ptr %177),
            cvt u64 i32 (mul i32 (regread i32 %190, constval i32 2)))
          */
        // We just assume its sth like register +/- sth patterns
        std::vector<VarID> dump_reg;
        std::vector<IntVal> dump_const;
        recursion(addr_expr->Opnd(0), dump_reg, dump_const);
        for (auto reg_tmp : dump_reg) {
          if (reg_tmp == register_number) {
            return prevStmt;
          }
        }
        ret_stmt = retLatest_Regassignment(prevStmt, register_number);
      } else {
        ret_stmt = retLatest_Regassignment(prevStmt, register_number);
      }
    } else {
      ret_stmt = retLatest_Regassignment(prevStmt, register_number);
    }
  }
  return ret_stmt;
}

void UDProfile::dep_iassign_expansion(IassignNode *iassign, set_check &dep) {
  BaseNode *rhs_expr = iassign->Opnd(1);
  if (rhs_expr->GetOpCode() == OP_regread) {
    // Case 1. regread u32 %13
    RegreadNode *regread = static_cast<RegreadNode *>(rhs_expr);
    dep.register_live.push(regread->GetRegIdx());
  } else if (rhs_expr->GetOpCode() == OP_constval) {
    // Case 2. constval i32 0 -> terminal
    ConstvalNode *constValNode = static_cast<ConstvalNode *>(rhs_expr);
    MIRConst *mirConst = constValNode->GetConstVal();
    if (mirConst != nullptr) {
      if (mirConst->GetKind() == kConstInt) {
        auto *const_to_get_value = safe_cast<MIRIntConst>(mirConst);
        dep.const_int64.push_back(const_to_get_value->GetValue());
      }
    }
  } else if (rhs_expr->GetOpCode() == OP_iread) {
    // Case 3. iread agg <* <$_TY_IDX334>> 0 (regread ptr %4) -> Only hold the ptr for deref
    std::vector<VarID> dump_reg;
    std::vector<IntVal> dump_const;
    recursion(rhs_expr, dump_reg, dump_const);
    for (VarID reg_temp : dump_reg) {
      dep.register_live.push(reg_temp);
    }
  } else {
    // Case 4. zext u32 8 (lshr u32 (regread u32 %4, constval i32 24))
    // Just assume it can be further expand and treat as a terminal...
    // Some of this of compound stmt are register
    // assigned by callassigned or function input register
    // Although there are some case didn't like this
    // We can set it as terminal register to prevent recursively deref
    // since it may crash
    // A proper SSA likely fix this issue
    std::vector<VarID> dump_reg;
    std::vector<IntVal> dump_const;
    recursion(rhs_expr, dump_reg, dump_const);
    for (VarID reg_temp : dump_reg) {
      dep.register_terminal.push_back(reg_temp);
    }
  }
}

void UDProfile::dep_constval_expansion(ConstvalNode *constValNode, set_check &dep) {
  MIRConst *mirConst = constValNode->GetConstVal();
  // we only trace int64
  // We didn't handle following cases
  // kConstFloatConst, MIRFloatConst
  // kConstDoubleConst, MIRDoubleConst
  if (mirConst != nullptr) {
    if (mirConst->GetKind() == kConstInt) {
      auto *const_to_get_value = safe_cast<MIRIntConst>(mirConst);
      dep.const_int64.push_back(const_to_get_value->GetValue());
    }
  }
}

void UDProfile::GetMatched(){
  float counter = 0 ; 
  float tot = this->StmtID2Check.size();
  for (const auto& pair0 : this->StmtID2Check){
    for (const auto& pair1 : this->StmtID2Check){
      if (pair0.first!=pair1.first){
        if (sat_check(pair0.second, pair1.second)){
          counter += 1;
        }
      }
    }
  }
  LogInfo::MapleLogger() << "Matched: " << counter/2.0 << " / " << ((tot*tot) - tot) /2.0 << "\n";
}

void UDProfile::dep_dassign_expansion(DassignNode *dassign, set_check &dep) {
  dep_expansion(dassign->GetRHS(), dep);
}

void UDProfile::dep_expansion(BaseNode *stmt, set_check &dep) {
  switch (stmt->GetOpCode()) {
    case OP_iassign: {
      IassignNode *iassign = static_cast<IassignNode *>(stmt);
      dep_iassign_expansion(iassign, dep);
      break;
    }
    case OP_regread: {
      RegreadNode *regread = static_cast<RegreadNode *>(stmt);
      dep.register_live.push(regread->GetRegIdx());
      break;
    }
    case OP_constval: {
      ConstvalNode *constValNode = static_cast<ConstvalNode *>(stmt);
      dep_constval_expansion(constValNode, dep);
      break;
    }
    case OP_conststr: {
      ConststrNode *conststr = static_cast<ConststrNode *>(stmt);
      dep.const_str.push_back(conststr->GetStrIdx());
      break;
    }
    case OP_conststr16: {
      Conststr16Node *conststr16 = static_cast<Conststr16Node *>(stmt);
      dep.const_str.push_back(conststr16->GetStrIdx());
      break;
    }
    case OP_dread: {
      DreadNode *dread = static_cast<DreadNode *>(stmt);
      dep.var_live.push(dread->GetStIdx().Idx());
      break;
    }
    case OP_addrof: {
      AddrofNode *addrof = static_cast<AddrofNode *>(stmt);
      dep.var_live.push(addrof->GetStIdx().Idx());
      break;
    }
    case OP_addroffunc: {
      // We don't handle function pointer
      break;
    }
    case OP_maydassign:
    case OP_dassign: {
      DassignNode *dassign = static_cast<DassignNode *>(stmt);
      dep.dassignID.push_back(dassign->GetStmtID());
      //LogInfo::MapleLogger() << "dassignID: " << int(dassign->GetStmtID()) << "\n";
      dep_dassign_expansion(dassign, dep);
      break;
    }
    case OP_dassignoff: {
      // TODO:
      // It is not documented in MAPLE IR.
      break;
    }
    default: {
      for (size_t i = 0; i < stmt->NumOpnds(); i++) {
        dep_expansion(stmt->Opnd(i), dep);
      }
      break;
    }
  }
}

StmtNode* UDProfile::retLatest_Varassignment(StmtNode *stmt, VarID var_number) {
  StmtNode *ret_stmt = nullptr;
  StmtNode *prevStmt = stmt->GetPrev();
  if (prevStmt != nullptr) {
    if (prevStmt->GetOpCode() == OP_dassign || prevStmt->GetOpCode() == OP_maydassign) {
      DassignNode *dassign = static_cast<DassignNode *>(prevStmt);
      // get the variable 
      if (var_number == dassign->GetStIdx().Idx()) {
        // get RHS 
        BaseNode *rhs_expr = dassign->GetRHS();
        std::vector<VarID> dump_reg;
        std::vector<IntVal> dump_const;
        recursion(rhs_expr, dump_reg, dump_const);
        // check if a = a + 1 condition exist 
        if (std::find(dump_reg.begin(), dump_reg.end(), var_number) != dump_reg.end()){
          ret_stmt = retLatest_Varassignment(prevStmt, var_number);
        } else{
          return prevStmt;
        }
      } else {
        ret_stmt = retLatest_Varassignment(prevStmt, var_number);
      }
    } else if (prevStmt->GetOpCode() == OP_iassign) {
      IassignNode *iassign = static_cast<IassignNode *>(prevStmt);
      BaseNode *addr_expr = iassign->Opnd(0);
      if (addr_expr->GetOpCode() == OP_dread) {
        DreadNode *dread = static_cast<DreadNode *>(addr_expr);
        if (var_number == dread->GetStIdx().Idx()) {
          return prevStmt;
        } else {
          ret_stmt = retLatest_Varassignment(prevStmt, var_number);
        }
      } else {
        ret_stmt = retLatest_Varassignment(prevStmt, var_number);
      }
    } else if (IsCallNode(prevStmt->GetOpCode())){
      // if is paratmeter, 
      // return the current 
      // treat function call as terminal 
      bool term = false;
      CallNode *callstmt = static_cast<CallNode *>(prevStmt);
      if (!callstmt->GetPragmas()->empty()) {
        // Same type ??? 
        // if the var number is in the pragma, then it is a terminal
        for (auto pragmaId : *callstmt->GetPragmas()) {
          if (pragmaId == var_number) {
            term = true;
            break;
          }
        }
      }
      if (term){
        return nullptr; 
      }else{
        ret_stmt = retLatest_Varassignment(prevStmt, var_number);
      }
      
    } else if (
            prevStmt->GetOpCode() == OP_if || 
            prevStmt->GetOpCode() == OP_brtrue || 
            prevStmt->GetOpCode() == OP_brfalse || 
            prevStmt->GetOpCode() == OP_dowhile || 
            prevStmt->GetOpCode() == OP_while ){
      return nullptr; 
    } else {
      ret_stmt = retLatest_Varassignment(prevStmt, var_number);
    }
  }
  return ret_stmt;
}

void UDProfile::gen_register_dep(StmtNode *stmt, set_check &br_tmp) {
  while (!br_tmp.var_live.empty()) {
    uint32_t var_to_check = br_tmp.var_live.top();
    auto iter = this->var_to_stmt.find(var_to_check);
    br_tmp.var_live.pop();
    if (iter != this->var_to_stmt.end()) {
      StmtNode *latest_stmt_tmp = retLatest_Varassignment(stmt, var_to_check);
      if (latest_stmt_tmp != nullptr) {
        set_check br_tmp_go_var;
        dep_expansion(latest_stmt_tmp, br_tmp_go_var );
        gen_register_dep(latest_stmt_tmp, br_tmp_go_var);
        br_tmp = commit(br_tmp, br_tmp_go_var);
      } else{
        br_tmp.var_terminal.push_back(var_to_check);
      }
    } else {
      br_tmp.var_terminal.push_back(var_to_check);
    }
  }
}

void UDProfile::recursion(BaseNode *stmt, std::vector<VarID> &stmt_var,
                                          std::vector<IntVal> &stmt_const ) {
  switch (stmt->GetOpCode()) {
    //case OP_regread: {
    case OP_dread: {
      DreadNode *dread = static_cast<DreadNode *>(stmt);
      stmt_var.push_back(dread->GetStIdx().Idx());
      break;
    }
    case OP_constval: {
      ConstvalNode *constValNode = static_cast<ConstvalNode *>(stmt);
      MIRConst *mirConst = constValNode->GetConstVal();
      if (mirConst != nullptr) {
        if (mirConst->GetKind() == kConstInt) {
          auto *const_to_get_value = safe_cast<MIRIntConst>(mirConst);
          stmt_const.push_back(const_to_get_value->GetValue());
        }
      }
      break;
    }
    default: {
      for (size_t i = 0; i < stmt->NumOpnds(); i++) {
        recursion(stmt->Opnd(i), stmt_var, stmt_const);
      }
    }
  }
}


bool UDProfile::sat_check(const set_check& a, const set_check& b) {
  if (compareVectors(a.var_terminal, b.var_terminal)
      /*
        compareVectors(a.const_int64,b.const_int64)
      */
  ) {
    return true;
  }
  return false;
}

bool UDProfile::dynamic_sat(const san_struct& a, const san_struct& b, bool SCSC) {
  // For SC-UC case, SC must be var a
  bool result = false;
  if (a.tot_ctr == 0) return result;
  if (b.tot_ctr == 0) return result;

  if (a.tot_ctr == b.tot_ctr ) {
    result = (a.false_ctr == b.false_ctr) || (a.false_ctr == b.true_ctr);
  } else if (!SCSC) {
    if ((a.tot_ctr == b.false_ctr || a.tot_ctr == b.true_ctr) && 
        (a.tot_ctr == a.false_ctr || a.tot_ctr == a.true_ctr)) {
      result = true;
    }
  }
  return result;
}


void UDProfile::Dump(){
  for (const auto& pair : this->StmtID2Check){
    LogInfo::MapleLogger() << "StmtID: " << pair.first ;
    set_check_print_dep(pair.second);
  }
}

void UDProfile::initME(MeFunction &mefunc, 
                       std::set<StmtID> asanStmtIDSet
){
  this->mefunc = &mefunc;
  this->asanStmtIDSet = asanStmtIDSet;
  for (StmtNode &stmt : this->mefunc->GetMirFunc()->GetBody()->GetStmtNodes()) {
    GetUD(stmt);
  }
  // Prepare a map for dassignID to set_check
  CreateDassignIDMap();
}

void UDProfile::GetUD(StmtNode &stmt) {
  /*
  - In current version, we cannot handle loops and if statements correctly
  */
  // Since current ASAN is placed before lowering, this code can't be triggered
  // // OP_regassign -> <REG> = <EXPR>
  switch (stmt.GetOpCode()) {
    case OP_regassign: {
      RegassignNode *regAssign = static_cast<RegassignNode *>(&stmt);
      if (this->reg_to_stmt.count(regAssign->GetRegIdx()) == 0) {
        this->reg_order.insert(regAssign->GetRegIdx());
      }
      this->reg_to_stmt[regAssign->GetRegIdx()].push_back(&stmt);
      break;
    }
    case OP_dassign:
    case OP_maydassign: {
      // TODO ... 
      // If <field-id> is not 0, then the variable must be a structure, 
      // and the assignment only applies to the specified field.
      DassignNode *dassign = static_cast<DassignNode *>(&stmt);
      // uint32
      if (this->var_to_stmt.count(dassign->GetStIdx().Idx()) == 0) {
        this->var_order.insert(dassign->GetStIdx().Idx());
      }
      this->var_to_stmt[dassign->GetStIdx().Idx()].push_back(&stmt);
      break;
    }
    // Unsupported OPCODE:
    // 1. iassignoff <prim-type> <offset> (<addr-expr>, <rhs-expr>)
    // 2023-02-07: I added iassignoff as interestedMemoryAccess, the address is
    // calculated by `<addr-expr> + offset`. Hence, the instrumented code is
    // simply the same as iassign
    // 2. callassigned
    // 2023-02-07: I added callassigned to transform the returned variables' names
    // there are dassign OpCodes inside callassigned instruction
    case OP_callassigned:{
      // We currently skip it, the retVar_XXX variable should not be instrumented
      // A Sanrazor considered function as a terminal 
      /*
                callassigned <func-name> (<opnd0>, ..., <opndn>) {
                dassign <var-name0> <field-id0>
                dassign <var-name1> <field-id1>
                ...
                dassign <var-namen> <field-idn> }
            */
      break;
    }
    case OP_iassign: {
      // syntax: iassign <type> <field-id> (<addr-expr>, <rhs-expr>)
      // %addr-expr = <rhs-expr>
      BaseNode *addr_expr = stmt.Opnd(0);
      // addr_expr have 3 cases
      // iread u64 <* <$_TY_IDX111>> 22 (regread ptr %177)
      if (addr_expr->GetOpCode() == OP_iread) {
        std::vector<VarID> dump_var;
        std::vector<IntVal> dump_const;
        recursion(addr_expr, dump_var, dump_const);
        for (VarID reg_tmp : dump_var) {
          if (this->var_to_stmt.count(reg_tmp) == 0) {
            this->var_order.insert(reg_tmp);
          }
          this->var_to_stmt[reg_tmp].push_back(&stmt);
        }
      } 
      else if (addr_expr->GetOpCode() == OP_dread) {
        // dread i64 %asan_shadowBase
        DreadNode *dread = static_cast<DreadNode *>(addr_expr);
        if (this->var_to_stmt.count(dread->GetStIdx().Idx()) == 0) {
          this->var_order.insert(dread->GetStIdx().Idx());
        }
        this->var_to_stmt[dread->GetStIdx().Idx()].push_back(&stmt);
      } else if (IsCommutative(addr_expr->GetOpCode())) {
        std::vector<VarID> dump_var;
        std::vector<IntVal> dump_const;
        recursion(addr_expr->Opnd(0), dump_var, dump_const);
        for (VarID reg_tmp : dump_var) {
          if (this->var_to_stmt.count(reg_tmp) == 0) {
            this->var_order.insert(reg_tmp);
          }
          this->var_to_stmt[reg_tmp].push_back(&stmt);
        }
      } 
      break;
    }
    // Checks 
    case OP_brtrue:
    case OP_brfalse:{
      set_check br_tmp;
      dep_expansion(stmt.Opnd(0), br_tmp);
      gen_register_dep(&stmt, br_tmp);
      //auto iter = this->StmtID2Check.find(stmt.GetStmtID());
      //CHECK_FATAL(iter == this->StmtID2Check.end() , "There is a dup entry in this->StmtID2Check.");
      this->StmtID2Check[stmt.GetStmtID()] = br_tmp;
      break;
    }
    case OP_if:{
      set_check if_tmp;
      dep_expansion(stmt.Opnd(0), if_tmp);
      gen_register_dep(&stmt, if_tmp);
      //auto iter = this->StmtID2Check.find(stmt.GetStmtID());
      //CHECK_FATAL(iter == this->StmtID2Check.end() , "There is a dup entry in this->StmtID2Check.");
      this->StmtID2Check[stmt.GetStmtID()] = if_tmp;
      IfStmtNode *ifNode = static_cast<IfStmtNode*>(&stmt);
      BlockNode *ThenBlock = ifNode->GetThenPart();
      if (ThenBlock != nullptr) {
        StmtNode *thenStmt = ThenBlock->GetFirst();
        while (thenStmt != nullptr) {
          if (thenStmt->GetOpCode() != OP_comment) {
            GetUD(*thenStmt);
          }
          thenStmt = thenStmt->GetNext();
        }
      }
      BlockNode *ElseBlock = ifNode->GetElsePart();
      if (ElseBlock != nullptr) {
        StmtNode *elseStmt = ElseBlock->GetFirst();
        while (elseStmt != nullptr) {
          if (elseStmt->GetOpCode() != OP_comment) {
            GetUD(*elseStmt);
          }
          elseStmt = elseStmt->GetNext();
        }
      }
      break;
    }
    case OP_dowhile:
    case OP_while:{
      set_check while_tmp;
      dep_expansion(stmt.Opnd(0), while_tmp);
      gen_register_dep(&stmt, while_tmp);
      //auto iter = this->StmtID2Check.find(stmt.GetStmtID());
      //CHECK_FATAL(iter == this->StmtID2Check.end() , "There is a dup entry in this->StmtID2Check.");
      this->StmtID2Check[stmt.GetStmtID()] = while_tmp;
      WhileStmtNode *whileNode = static_cast<WhileStmtNode*>(&stmt);
      BlockNode *Wbody = whileNode->GetBody();
      if (Wbody != nullptr) {
        StmtNode *wstmt_body = Wbody->GetFirst();
        while (wstmt_body != nullptr){
          GetUD(*(wstmt_body));
          wstmt_body = wstmt_body->GetNext();
        }
      }
      break;
    }

    case OP_doloop:{
      set_check doloop_tmp;
      dep_expansion(stmt.Opnd(1), doloop_tmp);
      gen_register_dep(&stmt, doloop_tmp);
      //auto iter = this->StmtID2Check.find(stmt.GetStmtID());
      //CHECK_FATAL(iter == this->StmtID2Check.end() , "There is a dup entry in this->StmtID2Check.");
      this->StmtID2Check[stmt.GetStmtID()] = doloop_tmp;

      DoloopNode *doloopNode = static_cast<DoloopNode*>(&stmt);
      BlockNode *Dbody = doloopNode->GetDoBody();
      for(auto &StmtInBlock : Dbody->GetStmtNodes()){
        GetUD(StmtInBlock);
      }
      break;
    }

    case OP_foreachelem:{
      ForeachelemNode *foreachelemNode = static_cast<ForeachelemNode*>(&stmt);
      // There are no comparsion checks ??
      BlockNode *Fbody = foreachelemNode->GetLoopBody();
      if (Fbody != nullptr) {
        StmtNode *fstmt_body = Fbody->GetFirst();
        while (fstmt_body != nullptr){
          GetUD(*(fstmt_body));
          fstmt_body = fstmt_body->GetNext();
        }
      }
      break;
    }

    // op if while dowhile 
    default: {
      for (size_t i = 0; i < stmt.NumOpnds(); i++) {
        GetUD(static_cast<StmtNode &>(*(stmt.Opnd(i))));
      }
      break;
    }
  }
}


bool UDProfile::UnmatchUserCheck(StmtID UcStmtID, std::set<StmtID> asanStmtIDSet){
  set_check * UcPtr = GetSetCheck(UcStmtID);
  if (UcPtr == nullptr) {
    // We might turn it to assertion 
    return true;
  }
  for (const auto& CheckID : asanStmtIDSet){
    set_check* ScPtr = GetSetCheck(CheckID);
    if (ScPtr == nullptr) continue;
    if (sat_check(*UcPtr, *ScPtr)) return false;
  }
  return true;
}

}

#endif  // ENABLE_MAPLE_SAN